import os.path
import xml.etree.ElementTree as ET

# 原xml路径 ，python 不区分 \ 和 / 确保目录下只有xml文件，没有其他目录
xmlpath = r'D:\dataset\abc\Annotations/'  

# 转换后txt文件存放路径，不要创建到xmlpath最后一级目录下，可以平级创建
txtpath = r'D:\dataset\abc\labels/'  

# 数据集标签，不区分先后顺序和索引
class_names = ['death','room_num_6','room_num_5','room_num_4','room_num_3','room_num_2','room_num_1','boss','boss_door','next_door','object','monster','hero','card','shop','return','again']

files = []
if not os.path.exists(txtpath):
    os.makedirs(txtpath)

for root, dirs, files in os.walk(xmlpath):
    None

number = len(files)
print(number)
i = 0
while i < number:
 
    name = files[i][0:-4]
    xml_name = name + ".xml"
    txt_name = name + ".txt"
    xml_file_name = xmlpath + xml_name
    txt_file_name = txtpath + txt_name
 
    xml_file = open(xml_file_name, 'r', encoding='utf-8')
    tree = ET.parse(xml_file)
    root = tree.getroot()
    # filename = root.find('name').text
 
    # image_name = root.find('filename').text
    w = int(root.find('size').find('width').text)
    h = int(root.find('size').find('height').text)
 
    f_txt = open(txt_file_name, 'w+')
    content = ""
 
    first = True
 
    for obj in root.iter('object'):
 
        name = obj.find('name').text
        class_num = class_names.index(name) 
        # class_num = 0
 
        xmlbox = obj.find('bndbox')
 
        x1 = float(xmlbox.find('xmin').text)
        x2 = float(xmlbox.find('xmax').text)
        y1 = float(xmlbox.find('ymin').text)
        y2 = float(xmlbox.find('ymax').text)
        
        x1 = round(x1)
        x2 = round(x2)
        y1 = round(y1)
        y2 = round(y2)
 
        if first:
            content += str(class_num) + " " + \
                       str((x1 + x2) / 2 / w) + " " + str((y1 + y2) / 2 / h) + " " + \
                       str((x2 - x1) / w) + " " + str((y2 - y1) / h)
            first = False
        else:
            content += "\n" + \
                       str(class_num) + " " + \
                       str((x1 + x2) / 2 / w) + " " + str((y1 + y2) / 2 / h) + " " + \
                       str((x2 - x1) / w) + " " + str((y2 - y1) / h)
 
    #print(str(i / (number - 1) * 100) + "%\n")
    #print(name)
    f_txt.write(content)
    f_txt.close()
    xml_file.close()
    i += 1
